# ==============================================================================
# Copyright (c) 2025 CompanyName. All rights reserved.
# Author:         22020873 陈泽欣
# Project:        Design of Deep Learning Fundamental Course
# Module:         grabscreen.py
# Date:           2025-05-24
# Description:    本模块用于实现屏幕图像的捕获与窗口管理功能，是整个骰子姿态识别系统中数据输入的第一环。
#                 主要功能包括：
#                 - 屏幕指定区域或全屏截图（grab_screen）；
#                 - 窗口初始化与置顶设置（setup_window）；
#                 支持从 Windows 桌面直接获取实时图像数据，为后续的图像处理和姿态估计提供原始图像输入。
# ==============================================================================

import cv2
import numpy as np
import win32gui, win32ui, win32con, win32api

# 初始化窗口
def setup_window(window_name='Mine-detect'):
    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
    cv2.resizeWindow(window_name, int(1920 // 2), int(1080 // 2))
    hwnd = win32gui.FindWindow(None, window_name)
    win32gui.SetWindowPos(hwnd, win32con.HWND_TOPMOST, 0, 0, 0, 0,
                          win32con.SWP_NOMOVE | win32con.SWP_NOSIZE)

def grab_screen(region=None):
    hwin = win32gui.GetDesktopWindow()

    if region:
        left, top, x2, y2 = region
        width = x2 - left + 1
        height = y2 - top + 1
    else:
        width = win32api.GetSystemMetrics(win32con.SM_CXVIRTUALSCREEN)
        height = win32api.GetSystemMetrics(win32con.SM_CYVIRTUALSCREEN)
        left = win32api.GetSystemMetrics(win32con.SM_XVIRTUALSCREEN)
        top = win32api.GetSystemMetrics(win32con.SM_YVIRTUALSCREEN)

    hwindc = win32gui.GetWindowDC(hwin)
    srcdc = win32ui.CreateDCFromHandle(hwindc)
    memdc = srcdc.CreateCompatibleDC()
    bmp = win32ui.CreateBitmap()
    bmp.CreateCompatibleBitmap(srcdc, width, height)
    memdc.SelectObject(bmp)
    memdc.BitBlt((0, 0), (width, height), srcdc, (left, top), win32con.SRCCOPY)

    signedIntsArray = bmp.GetBitmapBits(True)
    img = np.fromstring(signedIntsArray, dtype='uint8')
    img.shape = (height, width, 4)

    srcdc.DeleteDC()
    memdc.DeleteDC()
    win32gui.ReleaseDC(hwin, hwindc)
    win32gui.DeleteObject(bmp.GetHandle())

    return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
